from fileinput import filename
from numpy import extract
import pandas as pd
from collections import defaultdict
import glob
import json
import os
from tqdm import tqdm

NAME_LEN = 11

def convert_as_to_desed_labels(labels_str, label_map):
    """
    Convert AudioSet labels to DESED labels.
    """
    if pd.isna(labels_str):
        return ""
    labels = labels_str.split(",")
    desed_labels = []
    for label in labels:
        if label in label_map:
            desed_labels.append(label_map[label])
        else:
            desed_labels.append("other")
    return ",".join(set(desed_labels))

audioset_csvs = sorted(glob.glob("../../src/AudioSetStrong/meta/*.tsv"))
mapping_file = "../../src/DESED_AS_label_map.json"
desed_filenames = "../../src/DESED_all_filenames.json"

with open(mapping_file, "r") as f:
    label_map = json.load(f)    

# reverse the label map
reverse_label_map = defaultdict(str)
for key, value in label_map.items():
    for v in value:
        reverse_label_map[v] = key

label_map = reverse_label_map
valid_labels = set(label_map.values())
# load all desed filenames
with open(desed_filenames, "r") as f:
    desed_filenames = set(json.load(f))
# load audioset csvs
extracted_df = pd.DataFrame(columns=["segment_id", "label", "desed_labels"])
for csv in audioset_csvs:
    print(f"Processing {csv}...")
    df = pd.read_csv(csv, sep="\t")
    print(df.head())
    # unify labels from different files
    file_name_labels = df.groupby("segment_id")["label"].apply(lambda x: ",".join(set(",".join(x).split(","))))
    file_name_labels = file_name_labels.reset_index()
    file_name_labels["desed_labels"] =file_name_labels["label"].apply(lambda x: convert_as_to_desed_labels(x, label_map))
    print(file_name_labels.head())
    df_valid = file_name_labels[file_name_labels["desed_labels"].apply(lambda x: any([label in valid_labels for label in x.split(",")]))]
    print(f"Total rows: {len(file_name_labels)}, Valid rows: {len(df_valid)}")
    extracted_df = pd.concat([extracted_df, df_valid[["segment_id", "label", "desed_labels"]]], ignore_index=True)
    # print(df_valid.head())
# save the extracted dataframe to csv
os.makedirs("../../src/extracted/", exist_ok=True)
extracted_df.to_csv("../../src/extracted/DESED_AudiosetStrong_label_mapping.csv", index=False)

# find if desed file appeared in the csv
valid_labels = set()
extracted_df["desed_filename"] = extracted_df.apply(lambda row: f"{row['segment_id'][:NAME_LEN]}", axis=1)
as_filenames = set(extracted_df["desed_filename"].tolist())
desed_filenames = [x[1:1+NAME_LEN] for x in desed_filenames]  # remove the 'Y' prefix
intersec = as_filenames.intersection(desed_filenames)
if intersec:
    # print(f"Found {len(intersec)} desed filenames in {csv}.")
    # for filename in intersec:
    #     labels = extracted_df[extracted_df["desed_filename"] == filename]["desed_labels"].item()[1:-1]
    #     valid_labels.update(set(labels.split(",")))
    print("Duplicated files:", len(intersec))
else:
    print(f"No desed filenames found in {csv}.")
    